import pandas as pd
import numpy as np
from typing import Dict, Tuple


class Backtester:
    def __init__(self, initial_capital=200000):
        self.initial_capital = initial_capital
        self.results = None

    def run_backtest(self, data: pd.DataFrame, signals: pd.Series,
                     start_date=None, end_date=None) -> Dict:
        """
        执行回测
        :param data: 包含价格和信号的数据
        :param signals: 买卖信号（1=买入，-1=卖出）
        :param start_date: 回测开始日期
        :param end_date: 回测结束日期
        :return: 回测结果字典
        """
        # 确保索引对齐
        if not isinstance(signals, pd.Series):
            signals = pd.Series(signals, index=data.index)
        else:
            # 确保数据和信号索引完全一致
            signals = signals.reindex(data.index)

        # 检查是否有缺失值
        if signals.isnull().any():
            print("警告：存在未对齐的信号，将填充为0")
            signals = signals.fillna(0)

        # 筛选时间范围
        if start_date is not None and end_date is not None:
            # 确保都是 Timestamp 类型
            if isinstance(start_date, pd.Timestamp) and isinstance(end_date, pd.Timestamp):
                mask = (data.index >= start_date) & (data.index <= end_date)
                data = data.loc[mask].copy()
                signals = signals.loc[mask].copy()
            else:
                print("警告：日期类型不一致，跳过时间筛选")

        # 确保信号是Series类型
        if not isinstance(signals, pd.Series):
            signals = pd.Series(signals, index=data.index)

        # 初始化持仓
        positions = pd.DataFrame(index=data.index,
                                 columns=['holdings', 'cash', 'total', 'shares'])
        positions.iloc[0] = [0, self.initial_capital, self.initial_capital, 0]

        # 模拟交易
        positions = pd.DataFrame(index=data.index, columns=['holdings', 'cash', 'total', 'shares'])
        positions.iloc[0] = [0, self.initial_capital, self.initial_capital, 0]

        for i in range(1, len(data)):
            date = data.index[i]
            row = data.iloc[i]
            prev = positions.iloc[i - 1]
            current = positions.loc[date]

            # 执行买卖信号
            if signals.loc[date] == 1 and prev['shares'] == 0:  # 买入
                shares = prev['cash'] // row['Close']
                positions.at[date, 'shares'] = shares
                positions.at[date, 'cash'] = prev['cash'] - shares * row['Close']
            elif signals.loc[date] == -1 and prev['shares'] > 0:  # 卖出
                positions.at[date, 'cash'] = prev['cash'] + prev['shares'] * row['Close']
                positions.at[date, 'shares'] = 0
            else:  # 持仓不动
                positions.at[date, 'shares'] = prev['shares']
                positions.at[date, 'cash'] = prev['cash']

            # 计算当前资产
            positions.at[date, 'holdings'] = positions.at[date, 'shares'] * row['Close']
            positions.at[date, 'total'] = positions.at[date, 'cash'] + positions.at[date, 'holdings']

        # 计算绩效指标
        returns = positions['total'].pct_change()
        self.results = {
            'positions': positions,
            'returns': returns,
            'performance': self._calculate_performance(positions, returns)
        }
        return self.results

    def _calculate_performance(self, positions, returns) -> Dict:
        """计算关键绩效指标"""
        total_return = (positions['total'][-1] / self.initial_capital - 1) * 100
        sharpe_ratio = np.sqrt(252) * returns.mean() / returns.std()
        max_drawdown = (positions['total'].cummax() - positions['total']).max()

        return {
            'total_return (%)': round(total_return, 2),
            'annualized_return (%)': round(total_return / len(returns) * 252, 2),
            'sharpe_ratio': round(sharpe_ratio, 2),
            'max_drawdown': round(max_drawdown, 2),
            'win_rate': round((returns > 0).mean() * 100, 2)
        }
